package net.keebook.webapp.util;

import java.io.UnsupportedEncodingException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;

import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;

import org.bouncycastle.crypto.BufferedBlockCipher;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.engines.AESEngine;
import org.bouncycastle.crypto.modes.OFBBlockCipher;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.util.encoders.Base64;

public class AesCrypto {
	public static String encrypt(String s, String key) throws InvalidCipherTextException, NoSuchAlgorithmException, InvalidKeySpecException, UnsupportedEncodingException {
		// Prepare IV and Message
		byte[] data = s.getBytes("UTF-8");
		SecureRandom random = new SecureRandom();
		byte[] iv = random.generateSeed(16);
		
		// Generate Key
		SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
		KeySpec spec = new PBEKeySpec(key.toCharArray(), iv, 1, 256);
		SecretKey secret = factory.generateSecret(spec);

		// Initialize 
		BufferedBlockCipher aes = new BufferedBlockCipher(new OFBBlockCipher(new AESEngine(), 128));
		CipherParameters ivAndKey = new ParametersWithIV(new KeyParameter(secret.getEncoded()), iv);
		
		// Decrypt
		aes.init(true, ivAndKey);
		int minSize = aes.getOutputSize(data.length);
	    byte[] outBuf = new byte[minSize];
	    int length1 = aes.processBytes(data, 0, data.length, outBuf, 0);
	    int length2 = aes.doFinal(outBuf, length1);
	    int actualLength = length1 + length2;
	    byte[] result = new byte[actualLength];
	    System.arraycopy(outBuf, 0, result, 0, result.length);
	    
	    // Combine IV and encrypted message
	    byte[] finalResult = new byte[16 + result.length];
	    System.arraycopy(iv, 0, finalResult, 0, 16);
	    System.arraycopy(result, 0, finalResult, 16, result.length);
	    
	    return new String(Base64.encode(finalResult));
	}
	
	public static String decrypt(String s, String key) throws InvalidCipherTextException, NoSuchAlgorithmException, InvalidKeySpecException {
		// Extract IV and Message
		byte[] source = Base64.decode(s);
		byte[] data = new byte[source.length-16];
		byte[] iv = new byte[16];
		System.arraycopy(source, 0, iv, 0, 16);
		System.arraycopy(source, 16, data, 0, source.length-16);
		
		// Generate Key
		SecretKeyFactory factory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
		KeySpec spec = new PBEKeySpec(key.toCharArray(), iv, 1, 256);
		SecretKey secret = factory.generateSecret(spec);

		// Initialize 
		BufferedBlockCipher aes = new BufferedBlockCipher(new OFBBlockCipher(new AESEngine(), 128));
		CipherParameters ivAndKey = new ParametersWithIV(new KeyParameter(secret.getEncoded()), iv);
		
		// Decrypt
		aes.init(false, ivAndKey);
		int minSize = aes.getOutputSize(data.length);
	    byte[] outBuf = new byte[minSize];
	    int length1 = aes.processBytes(data, 0, data.length, outBuf, 0);
	    int length2 = aes.doFinal(outBuf, length1);
	    int actualLength = length1 + length2;
	    byte[] result = new byte[actualLength];
	    System.arraycopy(outBuf, 0, result, 0, result.length);
	    return new String(result);
	}
}
